import torch
import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter import ttk
import cv2
from PIL import Image, ImageTk
from ultralytics import YOLO
import os


class YOLOv11App:
    def __init__(self, root):
        self.root = root
        self.root.title("YOLOv11 GUI")

        # 设置主窗口的初始大小
        self.root.geometry("800x600")  # 宽800px, 高600px

        # 设置窗口的最大/最小尺寸
        self.root.resizable(True, True)  # 可以调整窗口的大小

        # 创建框架以管理布局
        self.frame = ttk.Frame(self.root)
        self.frame.pack(fill=tk.BOTH, expand=True, padx=20, pady=20)

        # 选择模型按钮
        self.model_button = ttk.Button(self.frame, text="选择YOLOv11模型", command=self.load_model)
        self.model_button.pack(pady=20)

        # 创建标签和输入框，要求输入归一化的宽和高
        self.normalization_label = ttk.Label(self.frame, text="请输入归一化的宽度和高度：")
        self.normalization_label.pack(pady=5)

        self.width_label = ttk.Label(self.frame, text="宽度：")
        self.width_label.pack(pady=5)

        self.width_entry = ttk.Entry(self.frame)
        self.width_entry.pack(pady=5)

        self.height_label = ttk.Label(self.frame, text="高度：")
        self.height_label.pack(pady=5)

        self.height_entry = ttk.Entry(self.frame)
        self.height_entry.pack(pady=5)

        # 上传单个图片按钮
        self.upload_button = ttk.Button(self.frame, text="上传单个图片", command=self.upload_image)
        self.upload_button.pack(pady=20)

        # 批量处理按钮
        self.batch_button = ttk.Button(self.frame, text="批量上传图片", command=self.batch_process_images)
        self.batch_button.pack(pady=20)

        # 进度条
        self.progress_bar = ttk.Progressbar(self.frame, length=400, mode="determinate")
        self.progress_bar.pack(pady=10)

        # 文本输出框
        self.text_output = tk.Text(self.frame, height=10, width=70)
        self.text_output.pack(pady=10)

        # 保存结果按钮
        self.save_button = ttk.Button(self.frame, text="保存结果", command=self.save_result)
        self.save_button.pack(pady=20)

        # 初始化模型和图片
        self.model = None
        self.image_path = None
        self.result_image = None  # 用于存储处理后的结果图像
        self.width = 640  # 默认宽度
        self.height = 640  # 默认高度
        self.save_folder = None  # 用于批量处理时保存路径

    def load_model(self):
        # 弹出文件对话框选择模型文件
        model_path = filedialog.askopenfilename(filetypes=[("PyTorch模型文件", "*.pt")])
        if model_path:
            try:
                # 使用 ultralytics 加载 YOLOv11 模型
                self.model = YOLO(model_path)
                messagebox.showinfo("加载成功", "模型加载成功！")
            except Exception as e:
                messagebox.showerror("加载失败", f"模型加载失败: {e}")

    def upload_image(self):
        if self.model is None:
            messagebox.showerror("错误", "请先加载模型！")
            return

        # 获取用户输入的归一化宽度和高度
        try:
            self.width = int(self.width_entry.get())
            self.height = int(self.height_entry.get())
        except ValueError:
            messagebox.showerror("错误", "请输入有效的数字作为宽度和高度！")
            return

        # 弹出文件对话框选择图片文件
        self.image_path = filedialog.askopenfilename(filetypes=[("图片文件", "*.jpg;*.jpeg;*.png")])
        if self.image_path:
            self.process_image(self.image_path)

    def process_image(self, image_path):
        # 读取原始图像
        img = cv2.imread(image_path)
        original_img = img.copy()  # 保持原始图像用于恢复

        # 获取图像的原始尺寸
        original_height, original_width = img.shape[:2]

        # 归一化并调整大小至用户指定的尺寸
        normalized_img = self.normalize_image(img)

        # 使用 YOLOv11 进行推理
        results = self.model(normalized_img)  # 传入归一化后的图片进行推理

        # 提取推理结果并渲染检测框
        img_result = results[0].plot()  # 使用 .plot() 绘制检测框，返回修改后的图片

        # 恢复图像至原始尺寸
        img_result = self.restore_image_size(img_result, original_width, original_height)

        # 保存处理后的图像
        self.result_image = img_result  # 保存处理后的结果图像

        # 弹出窗口显示图片
        self.show_image_in_popup(img_result)

    def normalize_image(self, img):
        # 使用 OpenCV 进行大小调整并归一化图片
        img_resized = cv2.resize(img, (self.width, self.height))  # 调整大小至用户指定的宽高
        img_normalized = img_resized / 255.0  # 归一化到 0-1 之间
        img_tensor = torch.from_numpy(img_normalized).float().permute(2, 0, 1).unsqueeze(0)  # 转为模型输入格式
        return img_tensor

    def restore_image_size(self, img_result, original_width, original_height):
        # 使用 OpenCV 恢复图片到原始尺寸
        img_result_resized = cv2.resize(img_result, (original_width, original_height))
        return img_result_resized

    def show_image_in_popup(self, img):
        # 创建新的弹窗显示图片
        popup = tk.Toplevel(self.root)
        popup.title("检测结果")

        # 转换为 PIL 图像以便在 Tkinter 中显示
        img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

        # 获取图片的宽高
        img_width, img_height = img_pil.size

        # 根据图片的纵横比调整弹窗大小
        popup.geometry(f"{img_width}x{img_height}")  # 设置弹窗的尺寸为图片尺寸

        # 转换为Tkinter可以显示的格式
        img_tk = ImageTk.PhotoImage(img_pil)

        # 创建并显示标签
        img_label = tk.Label(popup, image=img_tk)
        img_label.pack()

        # 保持对图片的引用，避免被垃圾回收
        img_label.image = img_tk

    def save_result(self):
        if self.result_image is None:
            messagebox.showerror("错误", "没有可保存的结果图像！")
            return

        # 弹出文件保存对话框
        save_path = filedialog.asksaveasfilename(defaultextension=".png",
                                                 filetypes=[("PNG 图片", "*.png"), ("JPEG 图片", "*.jpg")])
        if save_path:
            # 保存处理后的图像
            cv2.imwrite(save_path, self.result_image)
            messagebox.showinfo("保存成功", f"图像已保存到 {save_path}")

    def batch_process_images(self):
        if self.model is None:
            messagebox.showerror("错误", "请先加载模型！")
            return

        # 获取用户输入的归一化宽度和高度
        try:
            self.width = int(self.width_entry.get())
            self.height = int(self.height_entry.get())
        except ValueError:
            messagebox.showerror("错误", "请输入有效的数字作为宽度和高度！")
            return

        messagebox.showinfo(message="请选择要保存结果的文件夹")

        # 提示用户选择保存结果的文件夹
        self.save_folder = filedialog.askdirectory(title="选择保存结果的文件夹")
        if not self.save_folder:
            return

        # 弹窗提示用户选择了哪个文件夹
        messagebox.showinfo("选择保存文件夹", f"保存结果的文件夹是：{self.save_folder}")

        messagebox.showinfo(message="请选择要需要批量处理的图片的文件夹")

        # 提示用户选择需要批量处理的图片文件夹
        folder_path = filedialog.askdirectory(title="选择需要批量处理的图片文件夹")
        if not folder_path:
            return

        # 弹窗提示用户选择了批量处理的文件夹
        messagebox.showinfo("选择图像文件夹", f"批量处理的图片文件夹是：{folder_path}")

        # 获取文件夹内的所有图像文件
        image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]

        if not image_files:
            messagebox.showerror("错误", "该文件夹中没有图片！")
            return

        # 更新进度条最大值
        self.progress_bar["maximum"] = len(image_files)

        # 遍历所有图像文件并进行处理
        for idx, image_file in enumerate(image_files):
            image_path = os.path.join(folder_path, image_file)
            self.process_image_batch(image_path, image_file, idx)
        else:
            messagebox.showinfo("完成", "批量处理完成！")

    def process_image_batch(self, image_path, image_file, idx):
        # 读取图像并处理
        img = cv2.imread(image_path)
        original_img = img.copy()  # 保持原始图像用于恢复

        # 获取图像的原始尺寸
        original_height, original_width = img.shape[:2]

        # 归一化并调整大小
        normalized_img = self.normalize_image(img)

        # 使用 YOLOv11 进行推理
        results = self.model(normalized_img)

        # 渲染检测框
        img_result = results[0].plot()

        # 恢复图像至原始尺寸
        img_result = self.restore_image_size(img_result, original_width, original_height)

        # 保存处理后的图像
        save_path = os.path.join(self.save_folder, f"result_{image_file}")
        cv2.imwrite(save_path, img_result)

        # 更新进度条和文本输出
        self.progress_bar["value"] = idx + 1
        self.root.update_idletasks()  # 更新进度条
        self.text_output.insert(tk.END, f"正在处理: {image_file}\n")
        self.text_output.yview(tk.END)  # 滚动到最新输出

        # # 处理完成后提示用户
        # if idx + 1 == len(os.listdir(self.save_folder)):
        #     pass


if __name__ == "__main__":
    root = tk.Tk()
    app = YOLOv11App(root)
    root.mainloop()
